package org.idea.web.socket.factory;

import lombok.extern.slf4j.Slf4j;
import org.idea.web.socket.manager.SocketManager;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;

import javax.annotation.Resource;
import java.security.Principal;

/**
 * @Author linhao
 * @Date created in 6:17 下午 2021/5/8
 */
@Component
@Slf4j
public class WebSocketDecoratorFactory implements WebSocketHandlerDecoratorFactory {

    @Resource
    private SocketManager socketManager;

    @Override
    public WebSocketHandler decorate(WebSocketHandler webSocketHandler) {
        return new WebSocketHandlerDecorator(webSocketHandler){
            @Override
            public void afterConnectionEstablished(WebSocketSession session) throws Exception {
                log.info("新的连接接通了");
                Principal principal = session.getPrincipal();
                if(principal!=null) {
                    log.info("key = {},存入",principal.getName());
                    socketManager.add(principal.getName(),session);
                }
                super.afterConnectionEstablished(session);
            }

            @Override
            public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
                log.info("连接断开了");
                Principal principal = session.getPrincipal();
                if(principal!=null) {
                    log.info("key = {},断开",principal.getName());
                    socketManager.remove(principal.getName());
                }
                super.afterConnectionClosed(session, closeStatus);
            }
        };
    }
}
